# =============================
# 测试 SAC 起飞控制策略的效果


import logging
from stable_baselines3 import SAC  # 导入训练使用的算法
from train_ground_stabilize import TakeoffEnvWrapper  # 导入环境封装类

# ========== 日志初始化 ==========
logging.basicConfig(
    level=logging.INFO,  # 设置日志等级为 INFO
    format="%(asctime)s - %(levelname)s - %(message)s"  # 日志输出格式
)
logger = logging.getLogger("TestTakeoffSAC")  # 获取一个专用的 logger 实例

# ========== 主测试函数 ==========
def main():
    model_path = "results/models/sac_takeoff_task3.zip"  # 指定训练好的模型文件路径

    logger.info(f"加载模型：{model_path}")
    model = SAC.load(model_path)  # 加载 SAC 模型
    env = TakeoffEnvWrapper(fast_mode=True)  # 初始化测试环境，使用快速模式

    success_count = 0  # 成功起飞的回合数计数器
    episode_rewards = []  # 每个回合的总奖励，用于后续统计

    # 测试总共执行 5 个回合
    for ep in range(5):
        try:
            obs = env.reset()  # 重置环境，获取初始状态
        except Exception as e:
            logger.error(f"第 {ep+1} 轮初始化失败: {e}")  # 捕获异常，输出错误日志
            continue  # 跳过当前回合，进行下一个

        done = False  # 是否结束当前回合
        total_reward = 0  # 累计当前回合获得的总奖励

        # 当前回合循环，直到 done=True
        while not done:
            action, _ = model.predict(obs, deterministic=True)  # 使用模型预测动作（不加随机性）
            obs, reward, done, _ = env.step(action)  # 执行动作，获取新状态、奖励、完成标志
            total_reward += reward  # 累加本回合的奖励

        env.render()  # 绘制飞行轨迹图
        env.save_reward_plot()  # 绘制当前回合奖励趋势图
        episode_rewards.append(total_reward)  # 记录本回合总奖励

        # 检查是否起飞成功（起落架收起 + 高度高于初始高度 50m）
        gear_down = float(env.state[5]) > 0.9  # 起落架是否放下
        altitude = float(env.state[12])  # 当前高度
        success = not gear_down and altitude > env.initial_alt + 50  # 起飞判定逻辑
        success_count += int(success)  # 成功起飞的回合数 +1（成功为 True → int → 1）

        # 打印每一回合结果
        logger.info(f"Episode {ep+1}: Reward = {total_reward:.2f}, Success = {success}")

    # 所有回合测试完成后统计结果
    if episode_rewards:
        avg_reward = sum(episode_rewards) / len(episode_rewards)  # 计算平均奖励
        success_rate = success_count / len(episode_rewards) * 100  # 成功起飞比例（百分比）
        logger.info(f"平均奖励: {avg_reward:.2f}, 成功起飞率: {success_rate:.1f}%")

    env.close()  # 关闭环境（可用于释放资源）

# ========== 程序入口 ==========
if __name__ == "__main__":
    main()
